from Bio import PDB
import numpy as np
import pickle
import torch
from torch.utils.data import Dataset, DataLoader

from .dynamic_quantizer_components.dynamic_quantizer import DynamicQuantizerLitModule


class ProteinResidueDataset(Dataset):
    def __init__(self, data):
        self.spatial_x = []
        self.seq_x = []
        for item in data:
            self.spatial_x.append(item[0])
            self.seq_x.append(item[1])
        
    def __len__(self):
        return len(self.spatial_x)
    
    def __getitem__(self, idx):
        return {
            'spatial_x': self.spatial_x[idx],
            'seq_x': self.seq_x[idx],
        }


def compute_unit_normal_vector(A, B, C):
    AB = B - A
    AC = C - A
    normal_vector = np.cross(AB, AC)
    norm = np.linalg.norm(normal_vector)
    unit_normal_vector = normal_vector / norm
    return unit_normal_vector


def compute_rotation_matrices(v1_array, v2=np.array([1, 0, 0])):
    """
    计算从向量 v1_array (N,3) 旋转到 v2 (3,) 的旋转矩阵 (N,3,3)
    """
    v = np.cross(v1_array, v2)  # 计算旋转轴 (N,3)
    c = np.einsum('ij,j->i', v1_array, v2)  # 计算余弦值 (N,)
    s = np.linalg.norm(v, axis=1)  # 计算正弦值 (N,)
    
    R = np.repeat(np.eye(3)[np.newaxis, :, :], len(v1_array), axis=0)  # 初始化单位矩阵 (N,3,3)
    
    nonzero_mask = s > 1e-8  # 避免数值误差引起的除零
    
    if np.any(nonzero_mask):
        # v_nz = v[nonzero_mask] / s[nonzero_mask, np.newaxis]  # 归一化旋转轴 (M,3)
        v_nz = v[nonzero_mask]
        vx = np.zeros((len(v_nz), 3, 3))
        vx[:, 0, 1], vx[:, 1, 0] = -v_nz[:, 2], v_nz[:, 2]
        vx[:, 0, 2], vx[:, 2, 0] = v_nz[:, 1], -v_nz[:, 1]
        vx[:, 1, 2], vx[:, 2, 1] = -v_nz[:, 0], v_nz[:, 0]
        
        vx2 = np.einsum('nij,njk->nik', vx, vx)
        factor = ((1 - c[nonzero_mask]) / (s[nonzero_mask] ** 2))[:, np.newaxis, np.newaxis]
        
        R[nonzero_mask] += vx + vx2 * factor
    
    return R


class DynamicQuantizer:
    def __init__(
        self,
        batch_size=512,
        num_workers=0,
        encoder_path=None,
        kmeans_path=None,
    ):
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = DynamicQuantizerLitModule.load_from_checkpoint(encoder_path).eval().to(self.device)
        self.kmeans = pickle.load(open(kmeans_path, 'rb'))

    def __call__(self, pdb_file):
        pdb_parser = PDB.PDBParser(QUIET=False, PERMISSIVE=True)
        structure = pdb_parser.get_structure("protein", pdb_file)

        features = np.zeros((len(structure[0].get_list()[0]), 500, 6))

        for model_id, model in enumerate(structure):
            for chain in model:
                for residue_id, residue in enumerate(chain):
                    C_coord = residue['C'].get_coord()
                    N_coord = residue['N'].get_coord()
                    CA_coord = residue['CA'].get_coord()
                    O_coord = residue['O'].get_coord()
                    unit_normal_vector = compute_unit_normal_vector(C_coord, N_coord, O_coord)

                    features[residue_id][model_id][:3] = CA_coord
                    features[residue_id][model_id][3:] = unit_normal_vector

        dataset = []
        for i in range(features.shape[0]):
            for j in range(features.shape[0]):
                if i != j:
                    residue_0_unit_normals = features[i, :, 3:]
                    residue_0_CA_coords = features[i, :, :3]
                    residue_1_unit_normals = features[j, :, 3:]
                    residue_1_CA_coords = features[j, :, :3]

                    Rs = compute_rotation_matrices(residue_0_unit_normals)

                    delta_coords = residue_1_CA_coords - residue_0_CA_coords

                    coord_transformed = np.einsum('nij,nj->ni', Rs, delta_coords)
                    vector_transformed = np.einsum('nij,nj->ni', Rs, residue_1_unit_normals)
                    seq_dist = abs(i - j)

                    dataset.append((torch.tensor(np.concatenate([coord_transformed, vector_transformed], axis=1), dtype=torch.float32), torch.tensor(seq_dist, dtype=torch.float32)))


        dataset = ProteinResidueDataset(dataset)
        dataloader = DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers
        )

        embedding = []
        for batch in dataloader:
            spatial_x, seq_x = batch['spatial_x'], batch['seq_x']
            spatial_x, seq_x = spatial_x.to(self.device), seq_x.to(self.device)
            with torch.no_grad():
                hat_x = self.model.net.encoder.forward(spatial_x=spatial_x, seq_x=seq_x)
            embedding.append(hat_x.cpu().numpy())
        embedding = np.concatenate(embedding, axis=0)

        prediction = self.kmeans.predict(embedding)
        prediction_id = 0
        results = np.zeros((features.shape[0], features.shape[0]), dtype=np.int32)
        for i in range(features.shape[0]):
            for j in range(features.shape[0]):
                if i != j:
                    results[i][j] = prediction[prediction_id]
                    prediction_id += 1
        return results.tolist()

        
if __name__ == '__main__':
    pass